{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch RNN实现文本分帖"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! 1. 数据准备\n",
    "\n",
    "# 定义两个list分别存放两个板块的帖子数据\n",
    "academy_titles = []\n",
    "job_titles     = []\n",
    "with open('academy_titles.txt', encoding='utf8') as f:\n",
    "    for l in f:  # 按行读取文件\n",
    "        academy_titles.append(l.strip( ))  # strip 方法用于去掉行尾空格\n",
    "        \n",
    "with open('job_titles.txt', encoding='utf8') as f:\n",
    "    for l in f:  # 按行读取文件\n",
    "        job_titles.append(l.strip())       # strip 方法用于去掉行尾空格"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1217\n"
     ]
    }
   ],
   "source": [
    "#! 2. 统计文本中出现的字符数量\n",
    "char_set = set()\n",
    "for title in academy_titles:\n",
    "    for ch in title:\n",
    "        char_set.add(ch)\n",
    "for title in job_titles:\n",
    "    for ch in title:\n",
    "        char_set.add(ch)\n",
    "print(len(char_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#！3. 使用One-Hot编码表示数据\n",
    "import torch\n",
    "char_list  = list(char_set)\n",
    "n_chars    = len(char_list) + 1      #! 加一个 UNK\n",
    "\n",
    "def title_to_tensor(title):          #! 将title变为一个integer的vector\n",
    "    tensor = torch.zeros(len(title), dtype=torch.long)\n",
    "    for li, ch in enumerate(title):\n",
    "        try:\n",
    "            ind = char_list.index(ch)\n",
    "        except ValueError:\n",
    "            ind    = n_chars - 1     #! UNK的idx为最后一个\n",
    "        tensor[li] = ind\n",
    "    return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! 4. 定义模型\n",
    "import torch.nn as nn\n",
    "\n",
    "class RNN(nn.Module):\n",
    "    def __init__(self, word_count, embedding_size, hidden_size, output_size):\n",
    "        super(RNN, self).__init__()\n",
    "        \n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = torch.nn.Embedding(word_count, embedding_size)\n",
    "        self.i2h       = nn.Linear(embedding_size + hidden_size, hidden_size)\n",
    "        self.i2o       = nn.Linear(embedding_size + hidden_size, output_size)\n",
    "        self.softmax   = nn.LogSoftmax(dim=1)\n",
    "\n",
    "    def forward(self, input_tensor, hidden):\n",
    "        word_vector = self.embedding(input_tensor)\n",
    "        combined    = torch.cat((word_vector, hidden), 1)\n",
    "        hidden      = self.i2h(combined)          #! 输入 + 上一次的hidden \n",
    "        output      = self.i2o(combined)          #! 输入 + 上一次的hidden\n",
    "        output      = self.softmax(output)\n",
    "        return output, hidden\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, self.hidden_size)\n",
    "\n",
    "embedding_size = 100\n",
    "n_hidden       = 128\n",
    "n_categories   = 2\n",
    "rnn            = RNN(n_chars, embedding_size, n_hidden, n_categories)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "求食品专业生化 tensor([  30,  649, 1031,   35,  826,  889,  421]) tensor([30])\n",
      "output:\n",
      " tensor([[-0.4630, -0.9927]], grad_fn=<LogSoftmaxBackward0>)\n",
      "size of hidden:\n",
      " torch.Size([1, 128])\n"
     ]
    }
   ],
   "source": [
    "#! 5. 运行模型\n",
    "hidden         = rnn.initHidden()\n",
    "input_tensor   = title_to_tensor(academy_titles[0])\n",
    "# print(academy_titles[0], input_tensor, input_tensor[0].unsqueeze(dim=0))\n",
    "output, hidden = rnn(input_tensor[0].unsqueeze(dim=0), hidden)\n",
    "print('output:\\n', output)\n",
    "# print('hidden:\\n', hidden)\n",
    "print('size of hidden:\\n', hidden.size())\n",
    "\n",
    "# RNN模型的运行规律: 每次输入一个字符，每次使用上一次的hidden\n",
    "def run_rnn(rnn, input_tensor):\n",
    "    hidden = rnn.initHidden()\n",
    "    for i in range(input_tensor.size()[0]):\n",
    "        output, hidden = rnn(input_tensor[i].unsqueeze(dim=0), hidden)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train data size:  667\n",
      "Test data size:  287\n"
     ]
    }
   ],
   "source": [
    "#! 6. 数据预处理\n",
    "all_data   = []\n",
    "categories = [\"考研考博\", \"招聘信息\"]\n",
    "\n",
    "for line in academy_titles:\n",
    "    all_data.append((title_to_tensor(line), torch.tensor([0], dtype=torch.long)))   # 数据 + 标签\n",
    "for line in job_titles:\n",
    "    all_data.append((title_to_tensor(line), torch.tensor([1], dtype=torch.long)))\n",
    "    \n",
    "\n",
    "# 划分训练集与测试集    \n",
    "import random\n",
    "random.shuffle(all_data)\n",
    "data_len    = len(all_data)\n",
    "split_ratio = 0.7\n",
    "train_data  = all_data[:int(data_len*split_ratio)]\n",
    "test_data   = all_data[int(data_len*split_ratio):]\n",
    "print(\"Train data size: \", len(train_data))\n",
    "print(\"Test data size: \", len(test_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 194.15it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 637.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9024390243902439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 205.15it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 652.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9059233449477352\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 190.00it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 641.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9059233449477352\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 212.90it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 650.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9059233449477352\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 201.79it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 652.70it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9059233449477352\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 205.01it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 609.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9024390243902439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 209.88it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 599.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9024390243902439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 206.89it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 663.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9024390243902439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 203.68it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 620.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9024390243902439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 667/667 [00:03<00:00, 207.09it/s]\n",
      "100%|██████████| 287/287 [00:00<00:00, 618.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "accuracy 0.9024390243902439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# 7. 训练模型\n",
    "def train(rnn, criterion, input_tensor, category_tensor):\n",
    "    rnn.zero_grad()\n",
    "    output = run_rnn(rnn, input_tensor)\n",
    "    loss   = criterion(output, category_tensor)\n",
    "    loss.backward()\n",
    "\n",
    "    # 根据梯度更新模型的参数\n",
    "    for p in rnn.parameters():\n",
    "        p.data.add_(p.grad.data, alpha=-0.005)\n",
    "\n",
    "    return output, loss.item()\n",
    "\n",
    "def evaluate(rnn, input_tensor):\n",
    "    with torch.no_grad():\n",
    "        hidden = rnn.initHidden()\n",
    "        output = run_rnn(rnn, input_tensor)\n",
    "        return output\n",
    "    \n",
    "from tqdm import tqdm\n",
    "epoch          = 10\n",
    "embedding_size = 200\n",
    "n_hidden       = 10\n",
    "n_categories   = 2\n",
    "learning_rate  = 0.005\n",
    "criterion      = nn.NLLLoss()\n",
    "loss_sum       = 0\n",
    "all_losses     = []\n",
    "plot_every     = 100\n",
    "accuracy_best  = 0\n",
    "for e in range(epoch):\n",
    "    for ind, (title_tensor, label) in enumerate(tqdm(train_data)):\n",
    "        output, loss = train(rnn, criterion, title_tensor, label)\n",
    "        loss_sum += loss\n",
    "        if ind % plot_every == 0:\n",
    "            all_losses.append(loss_sum / plot_every)\n",
    "            loss_sum = 0\n",
    "    c = 0\n",
    "    for title, category in tqdm(test_data):\n",
    "        output = evaluate(rnn, title)\n",
    "        topn, topi = output.topk(1)\n",
    "        if topi.item() == category[0].item():\n",
    "            c += 1\n",
    "    print('accuracy', c / len(test_data))\n",
    "    if c / len(test_data) > accuracy_best:\n",
    "        torch.save(rnn, 'rnn_model.pkl')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "【校招】今日头条后端开发工程师 招聘信息\n"
     ]
    }
   ],
   "source": [
    "# 8. 验证\n",
    "def get_category(rnn, title):\n",
    "    input_tensor = title_to_tensor(title)\n",
    "    with torch.no_grad():\n",
    "        output = run_rnn(rnn, input_tensor)\n",
    "        top_val, top_idx = output.topk(1)\n",
    "        return categories[top_idx.item()]\n",
    "\n",
    "print(\"【校招】今日头条后端开发工程师\", get_category(rnn, \"【校招】今日头条后端开发工程师\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
